"""
使用装饰器模式实现带有人工审核的工具调用

这个版本修复了各种错误，并简化了实现
"""
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import interrupt
from langgraph.prebuilt import create_react_agent
import json
import asyncio
from typing import Optional, Tuple, Dict, Any, Callable, Union, cast
from functools import wraps
from langchain_openai import ChatOpenAI
from langgraph.types import Command
from langchain_tavily import TavilySearch
from langchain_core.tools import BaseTool, tool
from langchain_core.runnables import RunnableConfig
import os
import inspect

def human_in_the_loop():
    """
    装饰器，为工具函数添加人工审核功能
    
    返回:
        装饰器函数
    """
    def decorator(func: Callable):
        # 获取原始函数的签名信息
        tool_name = func.__name__
        
        @tool(tool_name)  # 使用工具名称作为第一个参数
        @wraps(func)  # 保留原始函数的元数据
        def wrapper(**kwargs):
            # 创建中断请求
            interrupt_data = {
                "tool": tool_name,
                "args": kwargs,
                "message": f"准备使用{tool_name}工具: 参数={kwargs}\n是否允许继续？\n输入 'yes' 接受，'no' 拒绝，或 'edit' 修改参数"
            }
            
            # 发起中断，等待人工审核
            response = interrupt(interrupt_data)
            print(f"审核响应: {response}")
            
            # 处理人工响应
            if response["type"] == "accept":
                print(f"工具调用已批准，执行 {tool_name} 中...")
                try:
                    result = func(**kwargs)
                    return result
                except Exception as e:
                    error_msg = f"工具调用失败: {str(e)}"
                    print(error_msg)
                    return error_msg
            
            elif response["type"] == "edit":
                edited_args = response.get("args", {})
                print(f"使用修改后的参数: {edited_args}")
                try:
                    result = func(**edited_args)
                    return result
                except Exception as e:
                    error_msg = f"修改后的工具调用失败: {str(e)}"
                    print(error_msg)
                    return error_msg
            
            elif response["type"] == "reject":
                print("工具调用被拒绝")
                return "该工具被拒绝使用，请尝试其他方法或提供不需要工具的回答。"
            
            else:
                error_msg = f"不支持的响应类型: {response['type']}"
                print(error_msg)
                return error_msg
                
        return wrapper
    
    return decorator

# 创建LLM
model = ChatOpenAI(model='gpt-4o-mini')

# 使用装饰器添加人工审核的Tavily搜索工具
@human_in_the_loop()
def tavily_search(query: str, search_depth: str = "basic"):
    """使用Tavily进行网络搜索"""
    try:
        search = TavilySearch(api_key=os.environ.get("TAVILY_API_KEY"))
        search_results = search.invoke({
            "query": query, 
            "search_depth": search_depth,
            "max_results": 5
        })
        return json.dumps(search_results, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"Tavily API调用失败: {str(e)}"

def create_tavily_search_agent():
    """创建一个带有Tavily搜索工具的智能体"""
    # 创建检查点保存器
    checkpointer = InMemorySaver() 
    
    # 智能体配置
    config = {
        "configurable": {
            "thread_id": "1"
        }
    }
    
    # 创建智能体
    agent = create_react_agent(
        model=model,
        tools=[tavily_search],
        checkpointer=checkpointer
    )
    
    return agent, config

async def main():
    """主函数，异步创建并运行智能体"""
    # 创建智能体
    agent, config = create_tavily_search_agent()
    
    # 开始流程
    result = await agent.ainvoke(
        {"messages": [{"role": "system", "content": "你会使用工具来帮助用户。如果工具使用被拒绝，请提示用户。"},
                      {"role": "user", "content": "中美贸易战最新消息"}]},
        config=config
    )
    
    # 处理可能的中断
    while "__interrupt__" in result:
        interrupt_info = result["__interrupt__"][0].value
        print("\n" + "="*50)
        print(f"收到中断请求")
        print(f"{interrupt_info.get('message', '')}")
        print("="*50)
        user_input = input("请输入：")
        
        # 处理用户输入
        while True:
            if user_input.lower() == "yes":
                result = await agent.ainvoke(
                    Command(resume={"type": "accept"}),
                    config=config
                )
                break
            elif user_input.lower() == "no":
                result = await agent.ainvoke(
                    Command(resume={"type": "reject"}),
                    config=config
                )
                break
            elif user_input.lower() == "edit":
                print("请输入新的搜索内容: ", end="")
                new_query = input()
                result = await agent.ainvoke(
                    Command(resume={"type": "edit", "args": {"query": new_query}}),
                    config=config
                )
                break
            else:
                print("无效输入，请输入 'yes'、'no' 或 'edit': ", end="")
                user_input = input()
    
    # 打印最终结果
    if result.get("messages"):
        print("\n=== 最终回答 ===")
        print(result['messages'][-1].content)
    
    return result

if __name__ == "__main__":
    asyncio.run(main())
